package org.hadoop.flume;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.InputStream;
import java.util.concurrent.CountDownLatch;

public class ZkUpload {
    /**
     * 上传文件
     *
     * @param args
     * @throws Exception
     */
    public static void main(String[] args) throws Exception {
        if (args.length < 3) {
            System.err.println("用法： <zkServer> <path> <file>");
            return;
        }
        String zkHosts = args[0];
        String path = args[1];
        String file = args[2];
        CountDownLatch countDownLatch = new CountDownLatch(1);
        ZooKeeper zooKeeper = new ZooKeeper(zkHosts, 3000, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                System.out.println("连接" + zkHosts + "成功");
                countDownLatch.countDown();
            }
        });
        countDownLatch.await();
        try {
            //检查文件是否存在;
            Stat stat = zooKeeper.exists(path, null);
            if (stat != null) {
                throw new RuntimeException("节点已经存在！");
            }
            //开始上传文件
            InputStream in = new FileInputStream(file);
            //判断文件大小，最大是1M
            int size = in.available();
            if (size > (1024 * 1024)) {
                throw new RuntimeException("上传的文件大小不能超过1M");
            }
            ByteArrayOutputStream bout = new ByteArrayOutputStream();
            byte[] bs = new byte[1024 * 4];
            int len = 0;
            while ((len = in.read(bs)) != -1) {
                bout.write(bs, 0, len);
            }
            bout.close();
            bs = bout.toByteArray();
            path = zooKeeper.create(path, bs, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            System.out.println("节点创建成功：" + path);
        } finally {
            zooKeeper.close();
        }
    }
}